import asyncio
from datetime import datetime
from typing import List

from llm_直播弹幕 import ormOB
from llm_直播弹幕.DatabaseManager import DatabaseManager
from llm_直播弹幕.redis.RedisManager import RedisManager


class PersistenceService:
    def __init__(self, redis: RedisManager, db: DatabaseManager):
        self.redis = redis
        self.db = db
        self.batch_size = 200  # 每批处理消息数
        self.batch_timeout = 10.0  # 最大等待时间（秒）
        self._buffer = []  # 消息缓冲区
        self.last_id = "0-0"  # Stream最后处理ID

    async def start_consuming(self):
        """启动消费服务（含批量处理机制）"""
        # 并行执行数据拉取和批量提交
        await asyncio.gather(
            self._continuous_fetch(),
            self._continuous_flush(),
        )

    async def _continuous_flush(self):
        """持续处理缓冲区的循环"""
        while True:
            await self._auto_flush()  # 保持原有_auto_flush内部逻辑不变

    async def _continuous_fetch(self):
        """持续拉取消息的循环"""
        while True:
            await self._fetch_messages()

    async def _fetch_messages(self):
        """从Redis Stream拉取消息"""
        messages = await self.redis.client.xread(
            streams={"STREAM:DANMAKU": self.last_id},
            count=self.batch_size * 2,  # 预取两倍数据量
            block=2000  # 阻塞2秒
        )

        if messages:
            stream_messages = messages[0][1]
            # 转换消息格式并缓存
            self._buffer.extend(self._parse_messages(stream_messages))
            # 更新最后处理ID
            self.last_id = stream_messages[-1][0]

    async def _auto_flush(self):
        """智能超时控制"""
        while True:
            # 双条件触发
            done, pending = await asyncio.wait(
                [
                    asyncio.create_task(self._wait_buffer()),
                    asyncio.create_task(self._wait_timeout())
                ],
                return_when=asyncio.FIRST_COMPLETED
            )

            # 取消未完成的任务
            for t in pending:
                t.cancel()

            await self._safe_batch_insert()

    async def _wait_buffer(self):
        """等待缓冲区填满"""
        while len(self._buffer) < self.batch_size:
            await asyncio.sleep(0.1)

    async def _wait_timeout(self):
        """等待超时触发"""
        await asyncio.sleep(self.batch_timeout)

    async def _safe_batch_insert(self):
        """带事务保障的批处理"""
        if not self._buffer:
            return

        success_ids = []
        try:
            # 数据库插入
            objects = [self._create_danmaku_obj(data) for _, data in self._buffer]
            await self._batch_insert(objects)

            # 记录成功ID
            success_ids = [msg_id for msg_id, _ in self._buffer]
        except Exception as e:
            print(f"处理失败: {e}")
        finally:
            # 无论成功与否都删除消息
            if success_ids:
                await self._delete_messages(success_ids)

            # 更新游标（取最后成功ID）
            if success_ids:
                self.last_id = success_ids[-1]
            self._buffer.clear()

    async def _batch_insert(self, objects: List[ormOB.Danmaku]):
        """优化后的批量插入方法"""
        async with self.db.get_session() as session:
            try:
                # 使用add_all批量操作
                session.add_all(objects)

                # # 针对SQLite的批量插入优化
                # await session.execute(
                #     "PRAGMA synchronous = OFF;"  # 禁用同步写入
                #     "PRAGMA journal_mode = MEMORY;"
                # )

                await session.commit()
            except Exception as e:
                await session.rollback()
                raise

    def _parse_messages(self, messages):
        """转换原始消息格式"""
        return [
            (msg_id, {k: v for k, v in data.items()})
            for msg_id, data in messages
        ]

    def _create_danmaku_obj(self, data: dict) -> ormOB.Danmaku:
        """完整字段映射"""
        return ormOB.Danmaku(
            danmu_user=data['danmu_user'],
            content=data['content'],
            create_time=datetime.fromtimestamp(int(data['create_time']) / 1000),
            danmu_time=datetime.fromtimestamp(int(data['danmu_time'])),
            user_mid=data['uid'],
            medal_level=data.get('medal_level', -1),
            medal_name=data.get('medal_name', ''),
            medal_from_user=data.get('medal_from_user', ''),
            medal_from_uid=data.get('medal_from_uid', 0),
            honor_level=data.get('honor_level', -1),
            guard_level=data.get('guard_level', -1),
            room_real_id=data.get('room_real_id', 0)
        )

    def _create_onlineRankCount_obj(self, data: dict) -> ormOB.OnlineRankCount:
        return ormOB.OnlineRankCount(
            create_time=datetime.fromtimestamp(int(data['create_time']) / 1000),
            room_real_id=data.get('room_real_id', -1),
            online_count=data.get('online_count', -1)
        )

    async def _delete_messages(self, msg_ids: list):
        """批量删除消息（兼容Redis5.0）"""
        await self.redis.client.xdel("STREAM:DANMAKU", *msg_ids)